//
//  AVL树.swift
//  Swift-LeetCode
//
//  Created by 卢悦明 on 2024/2/19.
//

import UIKit
/**
 avl树是平衡二叉搜索树的一种
 */
class AVLTree<E: Equatable & Comparable>: BinarySearchTree<E> {
//    
    override func addAfterNode(_ node: TreeNode<E>) {
        var node: TreeNode<E>? = node
        while let temp = node {
            if isBalanced(temp) {
                // 更新高度
                updateNodeHeight(temp)
            } else {
                // 恢复平衡
                rebalance(temp)
                break
            }
            node = node?.parent
            
        }
    }
    override func afterRemove(_ node: TreeNode<E>) {
        var node: TreeNode<E>? = node
        while let temp = node {
            if isBalanced(temp) {
                // 更新高度
                updateNodeHeight(temp)
            } else {
                // 恢复平衡
                rebalance(temp)
            }
            node = node?.parent
        }
    }
    override func creatNode(_ elem: E, parent: TreeNode<E>?) -> TreeNode<E> {
        AVLTreeNode(elem, parent: parent)
    }
//    节点是否平衡
    private func isBalanced(_ node: TreeNode<E>) -> Bool {
        guard let node = node as? AVLTreeNode else {
            return false
        }
        return abs(node.balanceFactor()) <= 1
    }
//    更新高度
    private func updateNodeHeight(_ node: TreeNode<E>?) {
        guard let node = (node as? AVLTreeNode) else {
            return
        }
        node.updateHeight()
    }
//    恢复平衡高度最低的不平衡节点
    private func rebalance(_ grand: TreeNode<E>) {
        guard let grand = grand as? AVLTreeNode<E> else {
            return
        }
        let parent = grand.tallerChild()
        let node = parent.tallerChild()
        if parent.isLeftChild() {
            if node.isLeftChild() { // LL
                rotateRight(grand)
            } else { // LR
                rotateLeft(parent)
                rotateRight(grand)
            }
        } else {
            if node.isRightChild() { // RR
                rotateLeft(grand)
            } else { // RL
                rotateRight(parent)
                rotateLeft(grand)
            }
            
        }
    }
//    节点向左旋转
    private func rotateLeft(_ grand: TreeNode<E>) {
        let parent = grand.right
        let child = parent?.left
        grand.right = child
        parent?.left = grand
        
        afterRotate(parent, grand, child)
    }
//    节点向右旋转
    private func rotateRight(_ grand: TreeNode<E>) {
        let parent = grand.left
        let child = parent?.right
        grand.left = child
        parent?.right = grand
        
        afterRotate(parent, grand, child)
    }
    // 更新父节点+节点高度
    private func afterRotate(_ parent: TreeNode<E>?, _ grand: TreeNode<E>, _ child:TreeNode<E>?) {
        // 更新parent的父节点
        parent?.parent = grand.parent
        if grand.isLeftChild() {
            grand.parent?.left = parent
        } else if grand.isRightChild(){
            grand.parent?.right = parent
        } else { //grand 是跟节点
            rootNode = parent
        }
        // 更新chil的父节点
        if child != nil {
            child?.parent = grand
        }
        // 更新guand的父节点
        grand.parent = parent
        
        // 更新节点高度
        updateNodeHeight(grand)
        updateNodeHeight(parent)
    }
}

class AVLTreeNode<E: Equatable & Comparable>: TreeNode<E> {
    var height: Int
    init(_ elem: E, parent: TreeNode<E>?) {
        height = 1
        super.init(elem: elem, parent: parent)
    }
//     获取节点的平衡因子
    func balanceFactor() -> Int {
        let leftHeight = self.left == nil ? 0 : (self.left as! AVLTreeNode).height
        let rightHeight = self.right == nil ? 0 : (self.right as! AVLTreeNode).height
        return leftHeight - rightHeight
    }
    func updateHeight() {
        let leftHeight = self.left == nil ? 0 : (self.left as! AVLTreeNode).height
        let rightHeight = self.right == nil ? 0 : (self.right as! AVLTreeNode).height
        height = max(leftHeight, rightHeight) + 1
    }
//    获取max高的子树
    func tallerChild() -> AVLTreeNode<E> {
        let leftHeight = self.left == nil ? 0 : (self.left as! AVLTreeNode).height
        let rightHeight = self.right == nil ? 0 : (self.right as! AVLTreeNode).height
        if leftHeight > rightHeight{
            return self.left as! AVLTreeNode
        } else if rightHeight > leftHeight {
            return self.right as! AVLTreeNode
        } else {
            return isLeftChild() ? (self.left as! AVLTreeNode<E>) : (self.right as! AVLTreeNode<E>);
        }
    }
}
